{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 相关说明\n",
    "- 数据集采用的是ctpn进行切分， 手动标记数据1000张， 然后采用ctpn进行切割编码。\n",
    "- 所有主干网络，均采用keras-applications中的代码， 其中InceptionResNetV2代码有修改， 将其中的padding操作参数全部修改成了“same”\n",
    "- 一个采用了4个主干网络： InceptionResNetV2， Xception， ResNet50， DenseNet169\n",
    "- 数据集采用了`sklearn.model_selection.KFold`划分成5份， 由于时间关系， 每个主干网络并没有训练5份， 而是训练了0,1，4一共3份， 还有一个原因是另外2份一直拟合的不好。\n",
    "- 识别部分采用的是 `cnn + ctc`， 训练集就是官方发布的39620张数据，然后采用ctpn进行切割编码，其中删除了几十张较差的图片\n",
    "- 训练集图片裁剪后尺寸大小不一致， 全部以高度为80进行等比缩放， 统一尺寸处理到：512×80， 长度超过512的进行resize， 不足512的左右补黑\n",
    "- 单模型的校验集全对率99.2%附近\n",
    "- 预测部分： 对于一个模型， 我会根据图片左右补黑的不同尺寸， 进行预测5份结果。 \n",
    "- 每一个主干网络， 根据数据集的不同， 我训练了3个模型， 然后每个模型预测了5份结果， 我采用了4个不同的主干网络， 一次一共训练了4×3×5=60份不同的结果。\n",
    "- 结果投票， 对同一个图片， 预测到的60个结果， 60个结果的每一个字符， 都进行投票， 得票最多的字符， 就是最终字符"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import random\n",
    "import os\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from imgaug import augmenters as iaa\n",
    "from tqdm import tqdm\n",
    "import shutil\n",
    "from sklearn.model_selection import KFold\n",
    "from tensorflow.keras.applications.resnet50 import *\n",
    "from tensorflow.keras.applications.densenet import *\n",
    "from tensorflow.keras.applications.nasnet import *\n",
    "from tensorflow.keras.applications.inception_resnet_v2 import *\n",
    "from tensorflow.keras.applications.inception_v3 import *\n",
    "from tensorflow.keras.applications.xception import *\n",
    "from tensorflow.keras.applications.vgg16 import *\n",
    "from tensorflow.keras.utils import to_categorical\n",
    "from tensorflow.keras.activations import *\n",
    "from tensorflow.keras.models import *\n",
    "from tensorflow.keras.layers import *\n",
    "from tensorflow.keras.models import *\n",
    "from tensorflow.keras.optimizers import *\n",
    "from tensorflow.keras.regularizers import l2\n",
    "from tensorflow.keras.utils import *\n",
    "import tensorflow.keras.backend as K\n",
    "from tensorflow.keras.callbacks import *\n",
    "import tensorflow as tf \n",
    "print(tf.__version__)\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "config = tf.ConfigProto() \n",
    "config.gpu_options.per_process_gpu_memory_fraction = 0.9\n",
    "K.set_session(tf.Session(config=config))\n",
    "\n",
    "\n",
    "# 定义ctc\n",
    "def ctc_lambda_func(args):\n",
    "    y_pred, labels, input_length, label_length = args\n",
    "    y_pred = y_pred[:, 2:, :]\n",
    "    return K.ctc_batch_cost(labels, y_pred, input_length, label_length)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "characters = '0123456789ABCDEFGHIJKLMNOPQRSTUWXYZ'\n",
    "width, height, n_len, n_class = 512, 80, 10, len(characters)+1\n",
    "print(width, height, n_len, n_class)\n",
    "\n",
    "# 训练集和校验集\n",
    "train_val_images_path = './datasets/train/'\n",
    "train_val_labels_path = './datasets/train_id_label.csv'\n",
    "csv_data = pd.read_csv(train_val_labels_path) \n",
    "csv_data = np.array(csv_data)\n",
    "\n",
    "kf = KFold(n_splits=5)\n",
    "kf_id = 4\n",
    "\n",
    "for i, (train_index, test_index) in enumerate(kf.split(csv_data)):\n",
    "    assert kf_id<5\n",
    "    if i == kf_id:\n",
    "        train_csv = csv_data[train_index]\n",
    "        val_csv = csv_data[test_index]\n",
    "\n",
    "print(\"train numbers: \", len(train_csv))\n",
    "print(\"val numbers: \", len(val_csv))\n",
    "\n",
    "test_images_path = './datasets/test/'\n",
    "test_images = os.listdir(test_images_path)\n",
    "print(\"test numbers: \", len(test_images))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "l2_rate = 1e-5\n",
    "rnn_size = 256\n",
    "BASE_MODEL_NAME=\"DenseNet169\"\n",
    "\n",
    "def crnn_model(base_model_name=\"InceptionResNetV2\", mode=\"train\"):\n",
    "    if mode==\"train\":\n",
    "        drop = 0.5\n",
    "    else:\n",
    "        drop = 0\n",
    "    input_tensor = Input(shape=(width, height, 3))\n",
    "    \n",
    "    if base_model_name==\"InceptionResNetV2\":\n",
    "        base_model = InceptionResNetV2(include_top=False, input_tensor = input_tensor)\n",
    "        x = base_model.get_layer('activation_161').output\n",
    "        \n",
    "    if base_model_name==\"Xception\":\n",
    "        base_model = Xception(include_top=False, input_tensor = input_tensor)\n",
    "        x = base_model.get_layer('block13_sepconv2_bn').output\n",
    "    \n",
    "    if base_model_name==\"ResNet50\":\n",
    "        base_model = ResNet50(include_top=False, input_tensor = input_tensor)\n",
    "        x = base_model.get_layer('activation_39').output\n",
    "        \n",
    "#     if base_model_name==\"DenseNet121\":\n",
    "#         base_model = DenseNet121(include_top=False, input_tensor = input_tensor)\n",
    "#         x = base_model.get_layer('pool4_conv').output\n",
    "        \n",
    "    if base_model_name==\"DenseNet169\":\n",
    "        base_model = DenseNet169(include_top=False, input_tensor = input_tensor)\n",
    "        x = base_model.get_layer('pool4_conv').output\n",
    "        \n",
    "#     if base_model_name==\"InceptionV3\":\n",
    "#         base_model = InceptionV3(include_top=False, input_tensor = input_tensor)\n",
    "#         x = base_model.get_layer('activation_74').output\n",
    "        \n",
    "        \n",
    "    conv_shape = x.get_shape().as_list()\n",
    "    rnn_length = conv_shape[1]\n",
    "    rnn_length -= 2\n",
    "    x = Reshape(target_shape=(conv_shape[1], conv_shape[2]*conv_shape[3]))(x)\n",
    "    x = Dropout(drop)(x)\n",
    "    x = Dense(rnn_size, kernel_regularizer=l2(l2_rate), bias_regularizer=l2(l2_rate))(x)\n",
    "    x = BatchNormalization(gamma_regularizer=l2(l2_rate), beta_regularizer=l2(l2_rate))(x)\n",
    "    x = Activation('relu')(x)\n",
    "\n",
    "    gru_1a = GRU(rnn_size, return_sequences=True, name='gru1a')(x)\n",
    "    gru_1b = GRU(rnn_size, return_sequences=True, go_backwards=True, name='gru1_b')(x)\n",
    "    gru1_merged = add([gru_1a, gru_1b])\n",
    "\n",
    "    gru_2a = GRU(rnn_size, return_sequences=True, name='gru2a')(gru1_merged)\n",
    "    gru_2b = GRU(rnn_size, return_sequences=True, go_backwards=True, name='gru2_b')(gru1_merged)\n",
    "    x = concatenate([gru_2a, gru_2b])\n",
    "    x = Dropout(drop)(x)\n",
    "\n",
    "    x = Dense(n_class, activation='softmax', kernel_regularizer=l2(l2_rate), bias_regularizer=l2(l2_rate))(x)\n",
    "    base_model = Model(inputs=input_tensor, outputs=x)\n",
    "\n",
    "    labels = Input(name='the_labels', shape=[n_len], dtype='float32')\n",
    "    input_length = Input(name='input_length', shape=(1,), dtype='int64')\n",
    "    label_length = Input(name='label_length', shape=(1,), dtype='int64')\n",
    "    loss_out = Lambda(ctc_lambda_func, name='ctc')([base_model.output, labels, input_length, label_length])\n",
    "\n",
    "    model = Model(inputs=(input_tensor, labels, input_length, label_length), outputs=loss_out)\n",
    "    return base_model, model, rnn_length\n",
    "\n",
    "K.clear_session()\n",
    "base_model, model, rnn_length = crnn_model(base_model_name=BASE_MODEL_NAME, mode=\"train\")\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "aug = iaa.Sequential([     \n",
    "    iaa.CropAndPad(\n",
    "        px=((0, 10), (0, 0), (0, 10), (0, 0)),\n",
    "        pad_mode=[\"constant\", \"edge\"],\n",
    "        pad_cval=(0, 128),\n",
    "    ),\n",
    "\n",
    "    # 仿射变换\n",
    "    iaa.Affine(\n",
    "        scale={\"x\": (0.9, 1.1), \"y\": (0.9, 1.1)},\n",
    "        rotate=(-5, 5),\n",
    "        cval=(0, 255),\n",
    "    ),\n",
    "    \n",
    "    iaa.Add((-10, 10), per_channel=0.5),\n",
    "    \n",
    "    # 增强和减弱图片的对比度\n",
    "    iaa.ContrastNormalization((0.8, 1.2)),\n",
    "    \n",
    "    # 图片更亮或者更暗\n",
    "    iaa.Multiply((0.8, 1.2), per_channel=0.5),\n",
    "    iaa.WithColorspace(\n",
    "         to_colorspace=\"HSV\",\n",
    "         from_colorspace=\"RGB\",\n",
    "         children=iaa.WithChannels(0, iaa.Add((0, 30)))\n",
    "     ),\n",
    "    \n",
    "     iaa.CoarseDropout(0.02, size_percent=0.5),\n",
    "    \n",
    "    ], random_order=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gen_data(csvs, train_val_images_path, mode='train', batch_size=64):\n",
    "    while True:\n",
    "        images = []\n",
    "        labels = []\n",
    "        #np.random.shuffle(val_csv)\n",
    "        for csv in csvs:\n",
    "            img = cv2.imread(os.path.join(train_val_images_path, csv[0]))\n",
    "            \n",
    "            if type(img)!=type(None):\n",
    "                scale = img.shape[1]/img.shape[0]\n",
    "                \n",
    "                new_width = int(height*scale)\n",
    "                if scale<3.5 or scale>7:\n",
    "                    continue\n",
    "                if new_width > width:\n",
    "                    new_width = width\n",
    "                img = cv2.resize(img, (new_width, height))\n",
    "                bg = np.zeros((height, width, 3))\n",
    "                #rand = random.randint(0, width-new_width)\n",
    "                rand = (width-new_width)//4 * random.randint(1, 4)\n",
    "                bg[:, rand:rand+new_width, :] = img\n",
    "                img = cv2.cvtColor(bg.astype('uint8'), cv2.COLOR_BGR2RGB)\n",
    "                lbl = [characters.find(x) for x in csv[1][1:]]\n",
    "                images.append(img)\n",
    "                labels.append(lbl)\n",
    "                \n",
    "                if len(images) == batch_size:    \n",
    "                    if mode=='train':\n",
    "                        images = aug(images=images)\n",
    "                    \n",
    "                    # images 预处理\n",
    "                    images = preprocess_input(np.array(images))\n",
    "                    \n",
    "                    images = images.transpose(0, 2, 1, 3)\n",
    "                    \n",
    "                    \n",
    "                    labels = np.array(labels)\n",
    "                    yield [images, labels, np.array([rnn_length]*batch_size), np.array([10]*batch_size)], np.ones(batch_size)\n",
    "                    images = []\n",
    "                    labels = []                \n",
    "\n",
    "g = gen_data(train_csv, train_val_images_path, mode='train', batch_size=3)\n",
    "[imgs, lbls, p, label_length_test], t = next(g)\n",
    "print(imgs.shape, lbls.shape, p, label_length_test, t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(batch_size=100, steps=20):\n",
    "    batch_acc = 0\n",
    "    generator = gen_data(val_csv, train_val_images_path, mode='val', batch_size=batch_size)\n",
    "    for i in range(steps):\n",
    "        [X_test, y_test, _, _], _  = next(generator)\n",
    "        y_pred = base_model.predict(X_test)\n",
    "        shape = y_pred[:,2:,:].shape\n",
    "        ctc_decode = K.ctc_decode(y_pred[:,2:,:], input_length=np.ones(shape[0])*shape[1])[0][0]\n",
    "        out = K.get_value(ctc_decode)[:, :n_len]\n",
    "\n",
    "        if out.shape[1] == n_len:\n",
    "            batch_acc += sum((y_test == out).all(axis=1))\n",
    "    return batch_acc / (batch_size*steps)\n",
    "\n",
    "class Evaluator(Callback):\n",
    "    def __init__(self):\n",
    "        self.accs = []\n",
    "    \n",
    "    def on_epoch_end(self, epoch, logs=None):\n",
    "        acc = evaluate(batch_size=100, steps=len(val_csv)//100)*100\n",
    "        self.accs.append(acc) \n",
    "        print('\\n acc: %f%%' % acc)\n",
    "evaluator = Evaluator()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 40\n",
    "model.compile(loss={'ctc': lambda y_true, y_pred: y_pred}, optimizer=Adam(lr=1e-4))\n",
    "reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=3, min_lr=1e-6, verbose=1)\n",
    "model_checkpoint = ModelCheckpoint('./models/%s_%d.h5'%(BASE_MODEL_NAME, kf_id), monitor='val_loss', save_best_only=True)\n",
    "earlystop = EarlyStopping(monitor='val_loss', patience=6, verbose=1, mode='auto')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#model.load_weights('./models/%s_%d.h5'%(BASE_MODEL_NAME, kf_id))\n",
    "history = model.fit_generator(\n",
    "    gen_data(train_csv, train_val_images_path, mode='train', batch_size=batch_size),\n",
    "    validation_data=gen_data(val_csv, train_val_images_path, mode='val', batch_size=batch_size), \n",
    "    validation_steps = len(val_csv)//batch_size,\n",
    "    steps_per_epoch=len(train_csv)//batch_size, \n",
    "    epochs=40, \n",
    "    callbacks=[reduce_lr, earlystop, model_checkpoint]) #evaluator,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "kf_ids = [0, 1, 4] \n",
    "BASE_MODEL_NAMES = [\"InceptionV3\"]\n",
    "        \n",
    "\n",
    "for BASE_MODEL_NAME in BASE_MODEL_NAMES:\n",
    "    for kf_id in kf_ids:\n",
    "        \n",
    "        K.clear_session()\n",
    "        base_model, model, rnn_length = crnn_model(base_model_name=BASE_MODEL_NAME, mode=\"test\")\n",
    "        model.load_weights('./models/%s_%d.h5'%(BASE_MODEL_NAME, kf_id))\n",
    "\n",
    "        for iii in range(1,6):\n",
    "            res = []\n",
    "            tensor = []\n",
    "\n",
    "            images = []\n",
    "            img_names = []\n",
    "            batch_size = 400\n",
    "            for img_name in tqdm(test_images):\n",
    "                img = cv2.imread(os.path.join(test_images_path, img_name))\n",
    "                scale = img.shape[1]/img.shape[0]\n",
    "                new_width = int(height*scale)\n",
    "                if new_width > width:\n",
    "                    new_width = width\n",
    "                img = cv2.resize(img, (new_width, height))\n",
    "                bg = np.zeros((height, width, 3))\n",
    "                rand = random.randint(0, width-new_width)\n",
    "                rand = (width-new_width)//6*iii\n",
    "                bg[:, rand:rand+new_width, :] = img\n",
    "                img = cv2.cvtColor(bg.astype('uint8'), cv2.COLOR_BGR2RGB).transpose(1, 0, 2)\n",
    "\n",
    "                images.append(img)\n",
    "                img_names.append(img_name)\n",
    "\n",
    "                if len(images)==batch_size:    \n",
    "                    images = preprocess_input(np.array(images))\n",
    "                    y_pred = base_model.predict(images)\n",
    "\n",
    "                    shape = y_pred[:,2:,:].shape\n",
    "                    ctc_decode = K.ctc_decode(y_pred[:,2:,:], input_length=np.ones(shape[0])*shape[1])[0][0]\n",
    "                    out = K.get_value(ctc_decode)[:, :n_len]\n",
    "\n",
    "                    for i in range(len(out)):\n",
    "                        label_list = [characters[int(idx)] for idx in out[i]]\n",
    "                        label = ''\n",
    "                        for c in label_list:\n",
    "                            label += c\n",
    "                        img_name = img_names[i]\n",
    "                        #print(img_name, label)\n",
    "                        res.append([img_name, label])\n",
    "\n",
    "                    images = []\n",
    "                    img_names = []\n",
    "\n",
    "            for i in range(len(res)):\n",
    "                res[i] = [res[i][0], ' '+res[i][1]]\n",
    "\n",
    "            pd_data = pd.DataFrame(np.array(res),columns=['name','label'])\n",
    "            print(\"BASE_MODEL_NAME\", BASE_MODEL_NAME, \"kf_id: \",kf_id, \"iii: \" ,iii)\n",
    "            pd_data.to_csv('./results/%s_%d_%d.csv'%(BASE_MODEL_NAME, kf_id, iii), index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "csv_list = os.listdir(\"./results/\")\n",
    "print(len(csv_list))\n",
    "csv_list.sort()\n",
    "for iii, sss in enumerate(csv_list):\n",
    "    print(iii, sss)\n",
    "#print(csv_list)\n",
    "csvs = []\n",
    "for csv in csv_list:\n",
    "    csv = np.array(pd.read_csv(\"./results/\" + csv))\n",
    "    csvs.append(csv)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "err_id = []\n",
    "def update_label(labels):\n",
    "    res = \"\"\n",
    "    for i in range(len(labels[0])):\n",
    "        cs = []\n",
    "        for s in labels:\n",
    "            if len(s)==11:\n",
    "                \n",
    "                if i>4 and (not (s[i] in \"0123456789\")):\n",
    "                    continue\n",
    "                cs.append(s[i])\n",
    "        d = {}\n",
    "        for c in cs:\n",
    "            if c in d:\n",
    "                d[c] += 1\n",
    "            else:\n",
    "                d[c] = 1\n",
    "        l = sorted(d.items(), key=lambda item: item[1], reverse=True) \n",
    "        res += l[0][0]\n",
    "    return res\n",
    "\n",
    "index = 0\n",
    "labels = []\n",
    "for i in tqdm(range(len(csvs[0]))):\n",
    "    lbls = []\n",
    "    for csv in csvs:\n",
    "        lbls.append(csv[i][1])\n",
    "    \n",
    "    index +=1\n",
    "    res = update_label(lbls)\n",
    "    \n",
    "    \n",
    "    if len(set(lbls)) >1: \n",
    "        print(index, \"-----------------------\")\n",
    "        print(lbls)\n",
    "        err_id.append(index)\n",
    "        \n",
    "        img = cv2.imread(\"./datasets/test/\" + csv[i][0])\n",
    "    labels.append(res)\n",
    "    \n",
    "    \n",
    "img_names = []\n",
    "for img_name in (test_images):\n",
    "    img_names.append(img_name)\n",
    "    \n",
    "res = []\n",
    "for i in range(len(img_names)):\n",
    "    res.append([img_names[i], labels[i]])\n",
    "    \n",
    "pd_data = pd.DataFrame(np.array(res),columns=['name','label'])\n",
    "pd_data.to_csv('./results/zzz-Xecp&&IncV2&&Dense&&resnet.csv', index=False)\n",
    "\n",
    "data = np.array(pd_data)\n",
    "\n",
    "print(\"==============================\")\n",
    "print(\"err_ids:\", len(err_id))\n",
    "for i,d in enumerate(data):\n",
    "    if i in err_id:\n",
    "        print(i, d)\n",
    "        shutil.copyfile(\"./datasets/test/\"+d[0],\"./datasets/err_test/\"+d[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
